import requests
import re
import os

class img_bdtb():
    def __init__(self,start_url,seeLZ):
        self.start_url =start_url
        self.seeLZ = '?see_lz=' + str(seeLZ)
        self.headers = {
            'User-Agent': "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.1 (KHTML, like Gecko) Chrome/22.0.1207.1 Safari/537.1"}
    def gethtml(self,pagenum):
        if self.seeLZ == 1:
            url = self.start_url + self.seeLZ + '&pn=' + str(pagenum)
        else:
            url = self.start_url + '?pn=' + str(pagenum)
        start_html = requests.get(url, headers=self.headers)
        return start_html.text
    def gettitle(self,page):
        page=self.gethtml(1)
        pattern=re.compile('<h1 class="core_title_txt.*?>(.*?)</h1>',re.S)
        result=re.search(pattern,page)
        if result:
            return result.group(1).strip()
        else:
            return None
    def getPageNum(self,page):
        page=self.gethtml(1)
        pattern=re.compile('<li class="l_reply_num.*?</span>.*?<span.*?>(.*?)</span>',re.S)
        result=re.search(pattern,page)
        if result:
            return result.group(1).strip()
        else:
            return None
    def getimg(self,page):
        pattern = re.compile('<div id="post_content_.*?<img class="BDE_Image.*?src="(.*?)".*?</div>',re.S)
        img_url=re.findall(pattern,page)
        imgcontent=[]
        for img in img_url:
            imgcontent.append(img)
        return  imgcontent
    def mkdir(self,path):
        path=path.strip()
        isExists = os.path.exists(os.path.join("D:\\img\\tieba", path))
        if not isExists:
            print(u'建了一个名字叫做', path, u'的文件夹！')
            os.makedirs(os.path.join("D:\\img\\tieba", path))
            os.chdir(os.path.join("D:\\img\\tieba", path)) ##切换到目录
            return True
        else:
            print(u'名字叫做', path, u'的文件夹已经存在了！')
            return False
    def save(self,imgcontent):
        for img in imgcontent:
            name=img[-9:-4]
            try:
                img = requests.get(img,headers=self.headers)
                f = open(name + '.jpg', 'ab')
                f.write(img.content)
                f.close()
            except FileNotFoundError:  ##捕获异常，继续往下走
                print(u'图片不存在已跳过：', img)
                return False
    def start(self):
        indexpage=self.gethtml(1)
        pagenum=self.getPageNum(indexpage)
        title=self.gettitle(indexpage)
        path=str(title).replace('?','_')
        self.mkdir(path)
        if pagenum==None:
            print("URL已失效，请重试")
            return
        try:
            print("该帖子共有" + str(pagenum) + "页")
            for i in range(1, int(pagenum) + 1):
                print("正在写入第" + str(i) + "页数据")
                page = self.gethtml(i)
                contents = self.getimg(page)
                self.save(contents)
        except IOError as e:
            print('写入异常，原因：',e.message)
        finally:
            print('写入完成!')

print('请输入帖子代号')
start_url = 'http://tieba.baidu.com/p/' + str(input(u'http://tieba.baidu.com/p/'))
seeLZ = input("是否只获取楼主发的图，是输入1，否输入0\n")
BDTB = img_bdtb(start_url, seeLZ)
BDTB.start()
